import tensorflow as tf
import numpy as np


# 下面的代码， 用 @tf.function 修饰，和不用 @tf.function 修饰，输出是不一样的!

# 原因：
# 当计算 f(a) 时，由于是第一次调用该函数，TensorFlow 进行了以下操作：
# 将函数内的代码依次运行了一遍（因此输出了文本）；
# 构建了计算图，然后运行了一次该计算图（因此输出了 1）。这里 tf.print(x) 可以作为计算图的节点，但 Python 内置的 print 则不能被转换成计算图的节点。因此，计算图中只包含了 tf.print(x) 这一操作；
# 将该计算图缓存到了一个哈希表中（如果之后再有类型为 tf.int32 ，shape 为空的张量输入，则重复使用已构建的计算图）。
# 计算 f(b) 时，由于 b 的类型与 a 相同，所以 TensorFlow 重复使用了之前已构建的计算图并运行（因此输出了 2）。这里由于并没有真正地逐行运行函数中的代码，所以函数第一行的文本输出代码没有运行。计算 f(b_) 时，TensorFlow 自动将 numpy 的数据结构转换成了 TensorFlow 中的张量，因此依然能够复用之前已构建的计算图。
# 计算 f(c) 时，虽然张量 c 的 shape 和 a 、 b 均相同，但类型为 tf.float32 ，因此 TensorFlow 重新运行了函数内代码（从而再次输出了文本）并建立了一个输入为 tf.float32 类型的计算图。
# 计算 f(d) 时，由于 d 和 c 的类型相同，所以 TensorFlow 复用了计算图，同理没有输出文本。

@tf.function
def f(x):
    print("The function is running in Python")
    tf.print(x)


a = tf.constant(1, dtype=tf.int32)
f(a)
b = tf.constant(2, dtype=tf.int32)
f(b)
b_ = np.array(2, dtype=np.int32)
f(b_)
c = tf.constant(0.1, dtype=tf.float32)
f(c)
d = tf.constant(0.2, dtype=tf.float32)
f(d)
